#include "gl\glut.h"
#include <Windows.h>
#include <iostream>
#include "io.h"
#include "core/lm.h"
#include "fft.h"

void SaveResult(int width, int height, char *filename, bool doFFT) {
	
	// align to byte
	
	// allocate memory for reading screen buffer
	unsigned char *pBuff = new unsigned char [width*height*3];
	// read buffer
	glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE, pBuff);
	// save to bmp
	SaveBMP(pBuff, width, height, filename);

	// if we want to perform FFT for the result
	if(doFFT) {

		// prepare value for FFT
		MYCOMPLEX *intensity = new MYCOMPLEX[width * height];
		for(int i = 0; i < width * height; i++) {
			intensity[i].real = double((int)pBuff[3*i] + (int)pBuff[3*i+1] + (int)pBuff[3*i+2]);
			intensity[i].imag = 0;
		}
		// perform FFT
		FFT2D(intensity, width, height, 1);
		
		// store spectrum in real
		for(int i = 0; i < width * height; i++) {
			intensity[i].real = sqrt(intensity[i].real * intensity[i].real + intensity[i].imag * intensity[i].imag);
		}

		// scale intensity.real to [0, 1]
		double imax = 0;
		for(int i = 0; i < width * height; i++) {
			imax = max(imax, intensity[i].real);
		}
		// avoid diveded by zero!
		if(imax > EPSILON(8))
			for(int i = 0; i < width * height; i++)
				intensity[i].real /= imax;

		// shift the origin
		// |1 2| --> |4 3|
		// |3 4|     |2 1|

		// origin [floor[M/2], floor[N/2]]
		int ox = width/2;
		int oy = height/2;

		// swap 1, 4 quadrant
		for(int w = 0; w < ox; w++)
			for(int h = 0; h < oy; h++)
				// (w,h) <--> (ox+w, oy+h)
				swap(intensity[w+h*width], intensity[ox+w + (oy+h)*width]);
		// swap 2, 3 quadrant
		for(int w = 0; w < ox; w++)
			for(int h = oy; h < height; h++)
				// (w,h) <--> (w+ox, h-oy)
				swap(intensity[w+h*width], intensity[ox+w +(h-oy)*width]);

		// convert to pixel value for writing bmp file
		unsigned char *fftImage = new unsigned char[width*height*3];
		for(int i = 0; i < width * height; i++) {
			fftImage[3*i] = fftImage[3*i+1] = fftImage[3*i+2] = TO_BYTE(log(1+intensity[i].real)/log(2.0));	// Log2 Transform + Gamma Correction
		}

		// find new name for bmp file
		// *.bmp --> *_fft.bmp
		char fftImageName[128];
		strcpy(fftImageName, filename);
		int length = strlen(fftImageName);
		fftImageName[length-1] = 't';
		fftImageName[length-2] = 'f';
		fftImageName[length-3] = 'f';
		fftImageName[length-4] = '_';
		strcat(fftImageName,".bmp\0");

		// save to bmp
		SaveBMP(fftImage, width, height, fftImageName);

		// free memory
		SAFE_DELETE_ARRAY(intensity);
		SAFE_DELETE_ARRAY(fftImage);
	}

	SAFE_DELETE_ARRAY(pBuff);
}

void SaveBMP(unsigned char* data, int width, int height, char *filename) {
	BITMAPFILEHEADER head;
	BITMAPINFOHEADER info;
	int	ptr, i, j;
	unsigned char pcrgb;

	FILE *fp=fopen(filename,"wb");
	if(fp==NULL)
	{
		printf("Can't open file %s\n",filename);
		return;
	}

	head.bfType	= 0x4d42;
	head.bfSize	= sizeof(BITMAPFILEHEADER)+sizeof(BITMAPINFOHEADER)+width*height*3; 
	head.bfReserved1= 0;
	head.bfReserved2= 0;
	head.bfOffBits	= sizeof(BITMAPFILEHEADER)+sizeof(BITMAPINFOHEADER);
	info.biSize	= sizeof(BITMAPINFOHEADER);
	info.biWidth	= width;
	info.biHeight	= height;
	info.biPlanes	= 1; 
	info.biBitCount	= 24;
	info.biCompression=0; 
	info.biSizeImage=0; 
	info.biXPelsPerMeter=0; 
	info.biYPelsPerMeter=0; 
	info.biClrUsed=0; 
	info.biClrImportant=0;
	fwrite(&head,sizeof(BITMAPFILEHEADER),1,fp);
	fwrite(&info,sizeof(BITMAPINFOHEADER),1,fp);

	ptr = 0;
	for( i = 0; i < height; i ++ )
	{
		for(j=0; j < width; j ++, ptr+= 3)
		{
			fwrite(&data[ptr+2],1,1,fp);
			fwrite(&data[ptr+1],1,1,fp);
			fwrite(&data[ptr],1,1,fp);
		}

		j = width*3;
		pcrgb=0;
		while((j%4) != 0)
		{	
			fwrite(&pcrgb,1,1,fp);
			j++;
		}
	}

	fclose(fp);

	Debug("%s was saved to bmp file\n", filename);
}